import paramiko
import zipfile
import os
import shutil
import getpass
import localconfig as CONFIG

def upload_file_to_server(local_file_path, remote_file_path, hostname, username, key_file_path, key_password=None):
    """
    Upload a local file to a remote server.
    :param local_file_path: The path of the local file.
    :param remote_file_path: The path where the file will be saved on the server.
    :param hostname: The hostname or IP address of the server.
    :param username: The username for the server.
    :param key_file_path: The path to the private key file.
    :param key_password: The password for the private key (optional).
    """
    # Create an SSH client
    ssh = paramiko.SSHClient()
    # Automatically add the host key
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    
    # Load the private key
    private_key = paramiko.Ed25519Key.from_private_key_file(key_file_path, password=key_password)
    
    # Connect to the server
    ssh.connect(hostname, username=username, pkey=private_key)

    # Create an SFTP client
    sftp = ssh.open_sftp()
    # Upload the file
    sftp.put(local_file_path, remote_file_path)
    # Close the SFTP client
    sftp.close()
    # Close the SSH client
    ssh.close()

def zip_folder(folder_path, zip_file_path):
    """
    Compress the specified folder and its contents into a ZIP file.
    :param folder_path: The path of the folder to be compressed.
    :param zip_file_path: The path of the generated ZIP file.
    """
    # Create a new ZIP file using the ZIP_DEFLATED compression algorithm
    with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        # Traverse all files in the folder and its subfolders
        for root, dirs, files in os.walk(folder_path):
            for file in files:
                # Get the file path
                file_path = os.path.join(root, file)
                # Get the relative path of the file
                relative_path = os.path.relpath(file_path, folder_path)
                # Add the file to the zip file
                zipf.write(file_path, relative_path)

# Write the main function
if __name__ == '__main__':
    # Call the function to compress the folder
    print('ziping...')
    if os.path.exists('learnctp.zip'):
        os.remove('learnctf.zip')
    if os.path.exists('learnctf'):
        shutil.rmtree('learnctf')
    shutil.copytree('.vitepress/dist', 'learnctf')
    zip_folder('learnctf', 'learnctf.zip')
    print('ziped.')
    # Call the function to upload the file
    kpass = getpass.getpass('key password:')
    print('uploading...')
    upload_file_to_server('learnctf.zip', f'/tmp/learnctf.zip', CONFIG.HOST, CONFIG.USER, CONFIG.KEY_PATH, kpass)
    print('uploaded.')
    # Remove
    print('removing...')
    os.remove('learnctf.zip')
    shutil.rmtree('learnctf')
    print('removed.')